// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/compute/api_vector.h>

#include <algorithm>
#include <memory>
#include <sstream>
#include <utility>
#include <vector>

#include <nebula/array/array_nested.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/compute/exec.h>
#include <nebula/compute/function.h>
#include <nebula/compute/function_internal.h>
#include <nebula/compute/kernels/vector_sort_internal.h>
#include <nebula/compute/registry.h>
#include <nebula/core/datum.h>
#include <nebula/core/record_batch.h>

#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>
#include <nebula/util/reflection_internal.h>

namespace nebula {

    namespace internal {

        using compute::DictionaryEncodeOptions;
        using compute::FilterOptions;
        using compute::NullPlacement;
        using compute::RankOptions;

        template<>
        struct EnumTraits<FilterOptions::NullSelectionBehavior>
                : BasicEnumTraits<FilterOptions::NullSelectionBehavior, FilterOptions::DROP,
                        FilterOptions::EMIT_NULL> {
            static std::string name() { return "FilterOptions::NullSelectionBehavior"; }

            static std::string value_name(FilterOptions::NullSelectionBehavior value) {
                switch (value) {
                    case FilterOptions::DROP:
                        return "DROP";
                    case FilterOptions::EMIT_NULL:
                        return "EMIT_NULL";
                }
                return "<INVALID>";
            }
        };

        template<>
        struct EnumTraits<DictionaryEncodeOptions::NullEncodingBehavior>
                : BasicEnumTraits<DictionaryEncodeOptions::NullEncodingBehavior,
                        DictionaryEncodeOptions::ENCODE, DictionaryEncodeOptions::MASK> {
            static std::string name() { return "DictionaryEncodeOptions::NullEncodingBehavior"; }

            static std::string value_name(DictionaryEncodeOptions::NullEncodingBehavior value) {
                switch (value) {
                    case DictionaryEncodeOptions::ENCODE:
                        return "ENCODE";
                    case DictionaryEncodeOptions::MASK:
                        return "MASK";
                }
                return "<INVALID>";
            }
        };

        template<>
        struct EnumTraits<NullPlacement>
                : BasicEnumTraits<NullPlacement, NullPlacement::AtStart, NullPlacement::AtEnd> {
            static std::string name() { return "NullPlacement"; }

            static std::string value_name(NullPlacement value) {
                switch (value) {
                    case NullPlacement::AtStart:
                        return "AtStart";
                    case NullPlacement::AtEnd:
                        return "AtEnd";
                }
                return "<INVALID>";
            }
        };

        template<>
        struct EnumTraits<RankOptions::Tiebreaker>
                : BasicEnumTraits<RankOptions::Tiebreaker, RankOptions::Min, RankOptions::Max,
                        RankOptions::First, RankOptions::Dense> {
            static std::string name() { return "Tiebreaker"; }

            static std::string value_name(RankOptions::Tiebreaker value) {
                switch (value) {
                    case RankOptions::Min:
                        return "Min";
                    case RankOptions::Max:
                        return "Max";
                    case RankOptions::First:
                        return "First";
                    case RankOptions::Dense:
                        return "Dense";
                }
                return "<INVALID>";
            }
        };

    }  // namespace internal

    namespace compute {

        // ----------------------------------------------------------------------
        // Function options

        namespace internal {
            namespace {
                using ::nebula::internal::DataMember;
                static auto kFilterOptionsType = get_function_options_type<FilterOptions>(
                        DataMember("null_selection_behavior", &FilterOptions::null_selection_behavior));
                static auto kTakeOptionsType = get_function_options_type<TakeOptions>(
                        DataMember("boundscheck", &TakeOptions::boundscheck));
                static auto kDictionaryEncodeOptionsType =
                        get_function_options_type<DictionaryEncodeOptions>(DataMember(
                                "null_encoding_behavior", &DictionaryEncodeOptions::null_encoding_behavior));
                static auto kRunEndEncodeOptionsType = get_function_options_type<RunEndEncodeOptions>(
                        DataMember("run_end_type", &RunEndEncodeOptions::run_end_type));
                static auto kArraySortOptionsType = get_function_options_type<ArraySortOptions>(
                        DataMember("order", &ArraySortOptions::order),
                        DataMember("null_placement", &ArraySortOptions::null_placement));
                static auto kSortOptionsType = get_function_options_type<SortOptions>(
                        DataMember("sort_keys", &SortOptions::sort_keys),
                        DataMember("null_placement", &SortOptions::null_placement));
                static auto kPartitionNthOptionsType = get_function_options_type<PartitionNthOptions>(
                        DataMember("pivot", &PartitionNthOptions::pivot),
                        DataMember("null_placement", &PartitionNthOptions::null_placement));
                static auto kSelectKOptionsType = get_function_options_type<SelectKOptions>(
                        DataMember("k", &SelectKOptions::k),
                        DataMember("sort_keys", &SelectKOptions::sort_keys));
                static auto kCumulativeOptionsType = get_function_options_type<CumulativeOptions>(
                        DataMember("start", &CumulativeOptions::start),
                        DataMember("skip_nulls", &CumulativeOptions::skip_nulls));
                static auto kRankOptionsType = get_function_options_type<RankOptions>(
                        DataMember("sort_keys", &RankOptions::sort_keys),
                        DataMember("null_placement", &RankOptions::null_placement),
                        DataMember("tiebreaker", &RankOptions::tiebreaker));
                static auto kPairwiseOptionsType = get_function_options_type<PairwiseOptions>(
                        DataMember("periods", &PairwiseOptions::periods));
                static auto kListFlattenOptionsType = get_function_options_type<ListFlattenOptions>(
                        DataMember("recursive", &ListFlattenOptions::recursive));
            }  // namespace
        }  // namespace internal

        FilterOptions::FilterOptions(NullSelectionBehavior null_selection)
                : FunctionOptions(internal::kFilterOptionsType),
                  null_selection_behavior(null_selection) {}

        constexpr char FilterOptions::kTypeName[];

        TakeOptions::TakeOptions(bool boundscheck)
                : FunctionOptions(internal::kTakeOptionsType), boundscheck(boundscheck) {}

        constexpr char TakeOptions::kTypeName[];

        DictionaryEncodeOptions::DictionaryEncodeOptions(NullEncodingBehavior null_encoding)
                : FunctionOptions(internal::kDictionaryEncodeOptionsType),
                  null_encoding_behavior(null_encoding) {}

        constexpr char DictionaryEncodeOptions::kTypeName[];

        RunEndEncodeOptions::RunEndEncodeOptions(std::shared_ptr<DataType> run_end_type)
                : FunctionOptions(internal::kRunEndEncodeOptionsType),
                  run_end_type{std::move(run_end_type)} {}

        ArraySortOptions::ArraySortOptions(SortOrder order, NullPlacement null_placement)
                : FunctionOptions(internal::kArraySortOptionsType),
                  order(order),
                  null_placement(null_placement) {}

        constexpr char ArraySortOptions::kTypeName[];

        SortOptions::SortOptions(std::vector<SortKey> sort_keys, NullPlacement null_placement)
                : FunctionOptions(internal::kSortOptionsType),
                  sort_keys(std::move(sort_keys)),
                  null_placement(null_placement) {}

        SortOptions::SortOptions(const Ordering &ordering)
                : FunctionOptions(internal::kSortOptionsType),
                  sort_keys(ordering.sort_keys()),
                  null_placement(ordering.null_placement()) {}

        constexpr char SortOptions::kTypeName[];

        PartitionNthOptions::PartitionNthOptions(int64_t pivot, NullPlacement null_placement)
                : FunctionOptions(internal::kPartitionNthOptionsType),
                  pivot(pivot),
                  null_placement(null_placement) {}

        constexpr char PartitionNthOptions::kTypeName[];

        SelectKOptions::SelectKOptions(int64_t k, std::vector<SortKey> sort_keys)
                : FunctionOptions(internal::kSelectKOptionsType),
                  k(k),
                  sort_keys(std::move(sort_keys)) {}

        constexpr char SelectKOptions::kTypeName[];

        CumulativeOptions::CumulativeOptions(bool skip_nulls)
                : FunctionOptions(internal::kCumulativeOptionsType), skip_nulls(skip_nulls) {}

        CumulativeOptions::CumulativeOptions(double start, bool skip_nulls)
                : CumulativeOptions(std::make_shared<Fp64Scalar>(start), skip_nulls) {}

        CumulativeOptions::CumulativeOptions(std::shared_ptr<Scalar> start, bool skip_nulls)
                : FunctionOptions(internal::kCumulativeOptionsType),
                  start(std::move(start)),
                  skip_nulls(skip_nulls) {}

        constexpr char CumulativeOptions::kTypeName[];

        RankOptions::RankOptions(std::vector<SortKey> sort_keys, NullPlacement null_placement,
                                 RankOptions::Tiebreaker tiebreaker)
                : FunctionOptions(internal::kRankOptionsType),
                  sort_keys(std::move(sort_keys)),
                  null_placement(null_placement),
                  tiebreaker(tiebreaker) {}

        constexpr char RankOptions::kTypeName[];

        PairwiseOptions::PairwiseOptions(int64_t periods)
                : FunctionOptions(internal::kPairwiseOptionsType), periods(periods) {}

        constexpr char PairwiseOptions::kTypeName[];

        ListFlattenOptions::ListFlattenOptions(bool recursive)
                : FunctionOptions(internal::kListFlattenOptionsType), recursive(recursive) {}

        constexpr char ListFlattenOptions::kTypeName[];

        namespace internal {
            void register_vector_options(FunctionRegistry *registry) {
                        KCHECK_OK(registry->add_function_options_type(kFilterOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kTakeOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kDictionaryEncodeOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kRunEndEncodeOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kArraySortOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kSortOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kPartitionNthOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kSelectKOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kCumulativeOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kRankOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kPairwiseOptionsType));
                        KCHECK_OK(registry->add_function_options_type(kListFlattenOptionsType));
            }
        }  // namespace internal

        // ----------------------------------------------------------------------
        // Direct exec interface to kernels

        turbo::Result<std::shared_ptr<Array>> NthToIndices(const Array &values,
                                                           const PartitionNthOptions &options,
                                                           ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum result, call_function("partition_nth_indices",
                                                           {Datum(values)}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> NthToIndices(const Array &values, int64_t n,
                                                           ExecContext *ctx) {
            PartitionNthOptions options(/*pivot=*/n);
            TURBO_MOVE_OR_RAISE(Datum result, call_function("partition_nth_indices",
                                                           {Datum(values)}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> SelectKUnstable(const Datum &datum,
                                                              const SelectKOptions &options,
                                                              ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum result,
                                call_function("select_k_unstable", {datum}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<Datum> ReplaceWithMask(const Datum &values, const Datum &mask,
                                             const Datum &replacements, ExecContext *ctx) {
            return call_function("replace_with_mask", {values, mask, replacements}, ctx);
        }

        turbo::Result<Datum> FillNullForward(const Datum &values, ExecContext *ctx) {
            return call_function("fill_null_forward", {values}, ctx);
        }

        turbo::Result<Datum> FillNullBackward(const Datum &values, ExecContext *ctx) {
            return call_function("fill_null_backward", {values}, ctx);
        }

        turbo::Result<std::shared_ptr<Array>> sort_indices(const Array &values,
                                                          const ArraySortOptions &options,
                                                          ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(
                    Datum result, call_function("array_sort_indices", {Datum(values)}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> sort_indices(const Array &values, SortOrder order,
                                                          ExecContext *ctx) {
            ArraySortOptions options(order);
            TURBO_MOVE_OR_RAISE(
                    Datum result, call_function("array_sort_indices", {Datum(values)}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> sort_indices(const ChunkedArray &chunked_array,
                                                          const ArraySortOptions &array_options,
                                                          ExecContext *ctx) {
            SortOptions options({SortKey("", array_options.order)}, array_options.null_placement);
            TURBO_MOVE_OR_RAISE(
                    Datum result, call_function("sort_indices", {Datum(chunked_array)}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> sort_indices(const ChunkedArray &chunked_array,
                                                          SortOrder order, ExecContext *ctx) {
            return sort_indices(chunked_array, ArraySortOptions(order), ctx);
        }

        turbo::Result<std::shared_ptr<Array>> sort_indices(const Datum &datum, const SortOptions &options,
                                                          ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum result,
                                call_function("sort_indices", {datum}, &options, ctx));
            return result.make_array();
        }

        turbo::Result<std::shared_ptr<Array>> Unique(const Datum &value, ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum result, call_function("unique", {value}, ctx));
            return result.make_array();
        }

        turbo::Result<Datum> DictionaryEncode(const Datum &value, const DictionaryEncodeOptions &options,
                                              ExecContext *ctx) {
            return call_function("dictionary_encode", {value}, &options, ctx);
        }

        turbo::Result<Datum> RunEndEncode(const Datum &value, const RunEndEncodeOptions &options,
                                          ExecContext *ctx) {
            return call_function("run_end_encode", {value}, &options, ctx);
        }

        turbo::Result<Datum> RunEndDecode(const Datum &value, ExecContext *ctx) {
            return call_function("run_end_decode", {value}, ctx);
        }

        const char kValuesFieldName[] = "values";
        const char kCountsFieldName[] = "counts";
        const int32_t kValuesFieldIndex = 0;
        const int32_t kCountsFieldIndex = 1;

        turbo::Result<std::shared_ptr<StructArray>> ValueCounts(const Datum &value, ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum result, call_function("value_counts", {value}, ctx));
            return turbo::checked_pointer_cast<StructArray>(result.make_array());
        }

        turbo::Result<std::shared_ptr<Array>> PairwiseDiff(const Array &array,
                                                           const PairwiseOptions &options,
                                                           bool check_overflow, ExecContext *ctx) {
            auto func_name = check_overflow ? "pairwise_diff_checked" : "pairwise_diff";
            TURBO_MOVE_OR_RAISE(Datum result,
                                call_function(func_name, {Datum(array)}, &options, ctx));
            return result.make_array();
        }

        // ----------------------------------------------------------------------
        // Filter- and take-related selection functions

        turbo::Result<Datum> Filter(const Datum &values, const Datum &filter,
                                    const FilterOptions &options, ExecContext *ctx) {
            // Invoke metafunction which deals with Datum kinds other than just Array,
            // ChunkedArray.
            return call_function("filter", {values, filter}, &options, ctx);
        }

        turbo::Result<Datum> Take(const Datum &values, const Datum &indices, const TakeOptions &options,
                                  ExecContext *ctx) {
            // Invoke metafunction which deals with Datum kinds other than just Array,
            // ChunkedArray.
            return call_function("take", {values, indices}, &options, ctx);
        }

        turbo::Result<std::shared_ptr<Array>> Take(const Array &values, const Array &indices,
                                                   const TakeOptions &options, ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum out, Take(Datum(values), Datum(indices), options, ctx));
            return out.make_array();
        }

        // ----------------------------------------------------------------------
        // Dropnull functions

        turbo::Result<Datum> DropNull(const Datum &values, ExecContext *ctx) {
            // Invoke metafunction which deals with Datum kinds other than just Array,
            // ChunkedArray.
            return call_function("drop_null", {values}, ctx);
        }

        turbo::Result<std::shared_ptr<Array>> DropNull(const Array &values, ExecContext *ctx) {
            TURBO_MOVE_OR_RAISE(Datum out, DropNull(Datum(values), ctx));
            return out.make_array();
        }

        // ----------------------------------------------------------------------
        // Cumulative functions

        turbo::Result<Datum> CumulativeSum(const Datum &values, const CumulativeOptions &options,
                                           bool check_overflow, ExecContext *ctx) {
            auto func_name = check_overflow ? "cumulative_sum_checked" : "cumulative_sum";
            return call_function(func_name, {Datum(values)}, &options, ctx);
        }

        turbo::Result<Datum> CumulativeProd(const Datum &values, const CumulativeOptions &options,
                                            bool check_overflow, ExecContext *ctx) {
            auto func_name = check_overflow ? "cumulative_prod_checked" : "cumulative_prod";
            return call_function(func_name, {Datum(values)}, &options, ctx);
        }

        turbo::Result<Datum> CumulativeMax(const Datum &values, const CumulativeOptions &options,
                                           ExecContext *ctx) {
            return call_function("cumulative_max", {Datum(values)}, &options, ctx);
        }

        turbo::Result<Datum> CumulativeMin(const Datum &values, const CumulativeOptions &options,
                                           ExecContext *ctx) {
            return call_function("cumulative_min", {Datum(values)}, &options, ctx);
        }

        turbo::Result<Datum> CumulativeMean(const Datum &values, const CumulativeOptions &options,
                                            ExecContext *ctx) {
            return call_function("cumulative_mean", {Datum(values)}, &options, ctx);
        }

    }  // namespace compute
}  // namespace nebula
